package com.byron.media.server.ffmpeg;

import lombok.extern.slf4j.Slf4j;
import org.bytedeco.javacpp.*;

import static org.bytedeco.javacpp.avcodec.*;
import static org.bytedeco.javacpp.avutil.*;
import static org.bytedeco.javacpp.swscale.*;

@Slf4j
public class H264toYuv {

    private FrameListener frameListener;

    private int width = 1280;
    private int height = 720;
    private int dpi = 1;

    private AVCodecContext inCodecContext = null;           // 解码工具类
    private avutil.AVFrame inFrame = null;					// 刚解码的frame数据，未转换
    private AVPacket inPacket;

    private AVCodecContext outCodecContext = null;  // 编码工具类
    private AVFrame outFrame = null;                // 转换大小之后的帧数据
    private AVPacket outPacket = null;              // 转换大小之后再编码的帧数据


    private swscale.SwsContext swsContext;

    public int handle(byte[] h264){
        int ret;
        inPacket.data().put(h264);
        inPacket.size(h264.length);
        ret = avcodec_send_packet(inCodecContext, inPacket);
        if (ret < 0) {
            inPacket.data().position(0);
            inPacket.size(0);
            log.info(String.format("Decode Error. %d ", ret));
            return -1;
        }

        ret = avcodec_receive_frame(inCodecContext, inFrame);

        if (ret < 0) {
            return -1;
        }

        // 解码出YUV
        int linesize0 = inFrame.linesize(0);
        int linesize1 = inFrame.linesize(1);
        int linesize2 = inFrame.linesize(2);
        int linesize3 = inFrame.linesize(3);
        int inWidth = inFrame.width();
        int inHeight = inFrame.height();
        int sampleRate = inFrame.sample_rate();

        // 转换尺寸
        if(swsContext == null){
            swsContext = sws_getContext(
                    inWidth, inHeight, inFrame.format(),
                    this.width, this.height, outFrame.format(),
//                        SWS_BICUBIC,
                    SWS_SINC,
                    null, null, (DoublePointer) null);
            if (swsContext == null) {
                log.info("Cannot initialize the conversion context");
                return -1;
            }
        }

        ret = sws_scale(swsContext, inFrame.data(), inFrame.linesize(),
                0, inHeight, outFrame.data(), outFrame.linesize());
        if(ret <= 0){
            return -1;
        }

        // 解码出YUV
        int outlinesize0 = outFrame.linesize(0);
        int outlinesize1 = outFrame.linesize(1);
        int outlinesize2 = outFrame.linesize(2);
        int outlinesize3 = outFrame.linesize(3);

        // y = linesize * height

        int dataLength0 = outlinesize0 * outFrame.height() ;
        byte[] data0 = new byte[dataLength0];
        outFrame.data(0).get(data0);

        int outinWidth = outFrame.width();
        int outinHeight = outFrame.height();
        int outsampleRate = outFrame.sample_rate();

        ret = avcodec_send_frame(outCodecContext, outFrame);
//        ret = avcodec_send_frame(outCodecContext, inFrame);
        if (ret == -22) {
            log.info(String.format("Encode Error. %d ", ret));
//                return -1;
        }
        outFrame.pts(outFrame.pts() + 1);
        ret = avcodec_receive_packet(outCodecContext, outPacket);
        if(ret < 0){
            log.info(String.format("Encode Error. %d ", ret));
            return -1;
        }

        if(frameListener != null){

        }
        return 0;
    }

    /**
     * 初始化H264
     */
    public int initFfmpeg(avformat.AVFormatContext inFormatContext){
//        avcodec_register_all();
//        avcodec_register_all();
        int video_stream_index = -1;
        int audio_stream_index = -1;
        for (int i = 0; i < inFormatContext.nb_streams(); ++i) {
            avformat.AVStream stream = inFormatContext.streams(i);
            log.info(String.format("type of the encoded data: %d", stream.codecpar().codec_id()));
            if (stream.codecpar().codec_type() == AVMEDIA_TYPE_VIDEO) {
                video_stream_index = i;
                log.info(String.format("dimensions of the video frame in pixels: width: %d, height: %d, pixel format: %d",
                        stream.codecpar().width(), stream.codecpar().height(), stream.codecpar().format()));
            } else if (stream.codecpar().codec_type() == AVMEDIA_TYPE_AUDIO) {
                audio_stream_index = i;
                log.info(String.format("audio sample format: %d", stream.codecpar().format()));
            }
        }

        if(video_stream_index == -1){
            return -1;
        }

        avformat.AVStream inStream = inFormatContext.streams(video_stream_index);

//        int inCodecId = AV_CODEC_ID_H264;
        int inCodecId = inStream.codecpar().codec_id();
        // 初始化解码器
        avcodec.AVCodec inCodec = avcodec_find_decoder(inCodecId);
        if (inCodec == null) {
            log.info("Codec not found");
            return -1;
        }

        // 初始化解码器的上下文
//        inCodecContext = avcodec_alloc_context3(inCodec);
        inCodecContext = inStream.codec();
        if (inCodecContext == null){
            log.info("Could not allocate video codec context");
            return -1;
        }

        // 打开解码器
        if (avcodec_open2(inCodecContext, inCodec, (avutil.AVDictionary) null) < 0) {
            log.info("Could not open codec");
            return -1;
        }

        inFrame = av_frame_alloc();
        inPacket = av_packet_alloc();
        av_new_packet(inPacket, 1024 * 1024);


        // 初始化编码器
//        int outCodecId = AV_CODEC_ID_H264;
        avcodec.AVCodec outCodec = avcodec_find_decoder(AV_CODEC_ID_H264);
        if (outCodec == null) {
            log.info("Codec not found");
            return -1;
        }

        // 初始化解码器的上下文
        outCodecContext = avcodec_alloc_context3(outCodec);
        if (outCodecContext == null){
            log.info("Could not allocate video codec context");
            return -1;
        }

        // 编码器参数
        outCodecContext
                //目标的码率，即采样的码率；显然，采样码率越大，视频大小越大
                .bit_rate(width * height * dpi)
                .rc_max_rate(width * height * dpi)
                .rc_min_rate(width * height * dpi)

                //编码目标的视频帧大小，以像素为单位
                .width(width)
//                .width(inCodecContext.width())
                .height(height)
//                .height(inCodecContext.height())

                //帧率的基本单位，我们用分数来表示，
                .time_base(new AVRational().num(1).den(10))

                //像素的格式，也就是说采用什么样的色彩空间来表明一个像素点
                .pix_fmt(AV_PIX_FMT_YUV420P)

                //每20帧插入1个I帧，I帧越少，视频越小
                .gop_size(20)

                //两个非B帧之间允许出现多少个B帧数
                //设置0表示不使用B帧
                //b 帧越多，图片越小
                .max_b_frames(0)

                //最大和最小量化系数
                .qmin(10)
                .qmax(50)

                //因为我们的量化系数q是在qmin和qmax之间浮动的，
                //qblur表示这种浮动变化的变化程度，取值范围0.0～1.0，取0表示不削减
                .qblur(0)

                //运动场景预判功能的力度，数值越大编码时间越长
                .me_pre_cmp(2)

                //DCT变换算法的设置，有7种设置，这个算法的设置是根据不同的CPU指令集来优化的取值范围在0-7之间
                .dct_algo(0)

                //这两个参数表示对过亮或过暗的场景作masking的力度，0表示不作
                .lumi_masking(0)
                .dark_masking(0);

//        AVDictionary param = new AVDictionary();
        // 新增语句，设置为编码延迟
        av_opt_set(outCodecContext.priv_data(), "preset", "superfast", 0);
        // 实时编码关键看这句，上面那条无所谓
        av_opt_set(outCodecContext.priv_data(), "tune", "zerolatency", 0);


        // 打开解码器
        int ret = avcodec_open2(outCodecContext, outCodec, (AVDictionary) null);
        if (ret < 0) {
            log.info("Could not open codec");
            return -1;
        }


        outFrame = av_frame_alloc();
        outPacket = av_packet_alloc();
        av_new_packet(outPacket, 1024 * 1024);

        outFrame
                .format(AV_PIX_FMT_YUV420P)
                .width(width)
                .height(height);
        av_image_alloc(
                outFrame.data(),
                outFrame.linesize(),
                outFrame.width(),
                outFrame.height(),
                AV_PIX_FMT_YUV420P,
                16);

        return 0;
    }


    /**
     * 释放资源
     * @return
     */
    public int releaseFfmpeg(){

        if(inCodecContext != null){
            avcodec_free_context(inCodecContext);
            inCodecContext.setNull();
            inCodecContext = null;
        }

        if(inFrame != null){
            av_frame_free(inFrame);
            inFrame.setNull();
            inFrame = null;
        }

        if(inPacket != null){
            av_packet_free(inPacket);
            inPacket.setNull();
            inPacket = null;
        }

        if(outCodecContext != null){
            avcodec_free_context(outCodecContext);
            outCodecContext.setNull();
            outCodecContext = null;
        }

        if(outFrame != null){
            av_frame_free(outFrame);
            outFrame.setNull();
            outFrame = null;
        }

        if(outPacket != null){
            av_packet_free(outPacket);
            outPacket.setNull();
            outPacket = null;
        }

        if (swsContext != null) {
            sws_freeContext(swsContext);
            swsContext.setNull();
            swsContext = null;
        }

        return 0;
    }

    public void setFrameListener(FrameListener frameListener) {
        this.frameListener = frameListener;
    }
}
